Average of levels in binary tree

Time: O(N); Space: O(H); easy

Given a non-empty binary tree, return the average value of the nodes on each level in the form of an array.

Example 1:

  3
 / \
9  20
  /  \
 15   7

Input: root = {TreeNode} [3,9,20,None,None,15,7]

Output: [3, 14.5, 11]

Explanation:

  • The average value of nodes:

    • on level 0 is 3,

    • on level 1 is 14.5,

    • and on level 2 is 11.

  • Hence return [3, 14.5, 11].

Constraints:

  • The range of node’s value is in the range of 32-bit signed integer.

[4]:

class TreeNode(object):
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

[5]:

class Solution1(object):
    """
    Time: O(N)
    Space: O(H)
    """
    def averageOfLevels(self, root):
        """
        :type root: TreeNode
        :rtype: List[float]
        """
        result = []
        q = [root]

        while q:
            total, count = 0, 0
            next_q = []

            for n in q:
                total += n.val
                count += 1
                if n.left:
                    next_q.append(n.left)
                if n.right:
                    next_q.append(n.right)
            q = next_q

            result.append(float(total) / count)

        return result

[6]:

s = Solution1()

root = TreeNode(3)
root.left, root.right = TreeNode(9), TreeNode(20)
root.right.left, root.right.right = TreeNode(15), TreeNode(7)
assert s.averageOfLevels(root) ==  [3, 14.5, 11]